from jili.core.printlog import print
import numpy as np
import pandas as pd

class cross:
    def __init__(self):
        self.fx=None
    def reset(self):
        self.fx=None
    def qs_size(self,x,field=""):
        if x[field]>0:
            return max(x["high"]-x["prelow"],x["high"]-x["low"])
        else:
            return max(x["prehigh"]-x["low"],x["high"]-x["low"])
    def qs_o_score(self,x,field=""):
        if x[field]>0:
            t=max(x["high"]-x["prelow"],x["high"]-x["low"])
            if t<=0:
                return 0
            else:
                return (x["high"]-x["close"])/t
        else:
            t=max(x["prehigh"]-x["low"],x["high"]-x["low"])
            if t<=0:
                return 0
            else:
                return (x["close"]-x["low"])/t
    def qs_o_score1(self,x,n=5,score="qs_o_score",size="qs_o_size"):
        if x[size]>=n:
            return x[score]
        else:
            return x[score]/2.5
    def qs_o_size(self,x,field=""):
        if x[field]>0:
            return x["high"]-x["close"]
        else:
            return x["close"]-x["low"]

    def bs(self, x, a):
        t=x[a]
        if pd.isna(t):
            return "h"
        else:
            if t > 0:
                return "b"
            elif t < 0:
                return "s"
            else:
                return "h"

    def cross01(self,x,a,b):
        if self.fx  is None:
            if x[a]>=x[b]:
                self.fx=1
                return 1
            else:
                self.fx=-1
                return -1
        else:
            if self.fx>0:
                if x[a]>=x[b]:
                    return self.fx
                else:
                    self.fx=-self.fx
                    return self.fx
            else:
                if x[a]<=x[b]:
                    return self.fx
                else:
                    self.fx=-self.fx
                    return self.fx
    def cross01_df(self,x,a,b):
        if np.isnan(x[a]):
            return np.nan
        if np.isnan(x[b]):
            return np.nan
        if self.fx  is None:
            if x[a]>=x[b]:
                self.fx=1
                return self.fx
            else:
                self.fx=-1
                return self.fx
        else:
            if self.fx>0:
                if x[a]>=x[b]:
                    return self.fx
                else:
                    self.fx=self.fx+1
                    self.fx=-self.fx
                    return self.fx
            else:
                if x[a]<=x[b]:
                    return self.fx
                else:
                    self.fx=self.fx-1
                    self.fx=-self.fx
                    return self.fx
    def cross(self,x,a,b):
        if self.fx  is None:
            if x[a]>=x[b]:
                self.fx=1
                return 0
            else:
                self.fx=-1
                return 0
        else:
            if self.fx==1:
                if x[a]>=x[b]:
                    return 0
                else:
                    self.fx=-1
                    return -1
            elif self.fx==-1:
                if x[a]<=x[b]:
                    return 0
                else:
                    self.fx=1
                    return 1
    def cross_df(self,x,a,b):
        if np.isnan(x[a]):
            return np.nan
        if np.isnan(x[b]):
            return np.nan
        if self.fx  is None:
            if x[a]>=x[b]:
                self.fx=1
                return self.fx
            else:
                self.fx=-1
                return self.fx
        else:
            if self.fx>0:
                if x[a]>=x[b]:
                    return np.nan
                else:
                    self.fx=self.fx+1
                    self.fx=-self.fx
                    return self.fx
            else:
                if x[a]<=x[b]:
                    return np.nan
                else:
                    self.fx=self.fx-1
                    self.fx=-self.fx
                    return self.fx
    def recoveryk(self,k,persist):
        ll=[]
        for i in k.columns:
            if i not in persist:
                ll.append(i)
        if ll:
            k.drop(ll, inplace=True, axis=1)
    def run_v_score(self,k,arg):
        a=arg["a"]
        b=arg["b"]
        k["b-a"]=k[a]-k[a]
        k["v_b-a"]=k["b-a"].diff(1)

    def run_score(self,k,arg):
        out={}
        self.reset()
        k["Cross01"] = k.apply(self.cross01_df, axis=1, a=arg["a"], b=arg["b"])
        self.reset()
        k["Cross"] = k.apply(self.cross_df, axis=1, a=arg["a"], b=arg["b"])
        t1 = k[pd.notna(k["Cross"])][["close", "Cross"]]
        t2 = k.groupby(by=["Cross01"])["high"].max()
        t3 = k.groupby(by=["Cross01"])["low"].min()
        t1["timekey"] = t1.index
        t1.index = t1["Cross"]
        t4 = pd.concat([t1, t2, t3], axis=1)
        t4.index = t4["timekey"]
        t4.sort_index(inplace=True)
        t4["prehigh"] = t4["high"] - t4["high"].diff(1)
        t4["prelow"] = t4["low"] - t4["low"].diff(1)
        testn=len(t4)
        t4.dropna(inplace=True)
        if len(t4)>0:
            t4["qs_o_size"] = t4.apply(self.qs_o_size, field="Cross", axis=1)
            t4["qssize"] = t4.apply(self.qs_size, field="Cross", axis=1)
            t4["qs_o_score"] = t4.apply(self.qs_o_score, field="Cross", axis=1)
            t4["qs_o5_score"] = t4.apply(self.qs_o_score1, n=5,score="qs_o_score",size="qs_o_size", axis=1)

            out["o_times"]=len(t4)
            out["o3_rate"] = len(t4[t4["qs_o_size"] >= 3]) / len(t4)
            out["qs3_rate"] = len(t4[t4["qssize"] >= 3]) / len(t4)
            out["o5_rate"]=len(t4[t4["qs_o_size"]>=5])/len(t4)
            out["qs5_rate"]=len(t4[t4["qssize"]>=5])/len(t4)
            out["o6_rate"] = len(t4[t4["qs_o_size"] >= 6]) / len(t4)
            out["qs6_rate"] = len(t4[t4["qssize"] >= 6]) / len(t4)
            out["o8_rate"] = len(t4[t4["qs_o_size"] >= 8]) / len(t4)
            out["qs8_rate"] = len(t4[t4["qssize"] >= 8]) / len(t4)
            out["o10_rate"] = len(t4[t4["qs_o_size"] >= 10]) / len(t4)
            out["qs10_rate"] = len(t4[t4["qssize"] >= 10]) / len(t4)
            out["score"]=t4["qs_o_score"].mean()
            out["score5"] = t4["qs_o5_score"].mean()
            self.recoveryk(t4,["qssize","qs_o_size","qs_o5_score","close"])
            a = t4.describe()
            out["describe"]=a.to_dict()
        else:
            print("无机会",testn)
            out["o_times"] = len(t4)
            out["o3_rate"] = 0
            out["qs3_rate"] = 0
            out["o5_rate"] = 0
            out["qs5_rate"] = 0
            out["o6_rate"] = 0
            out["qs6_rate"] = 0
            out["o8_rate"] = 0
            out["qs8_rate"] = 0
            out["o10_rate"] =0
            out["qs10_rate"] =0
            out["score"] = 0
            out["score5"] = 0
            out["describe"] = {}
        return out


if __name__ == '__main__':
#    from jili.calc.apply.cross import cross
    import pandas as pd
    from talib import abstract
    from jili.data.db import getk1mbar_pkl
    k = getk1mbar_pkl("rb2105")
    k = pd.DataFrame(k)
    sma = abstract.Function("SMA")
    sma.parameters
    sma.set_parameters({"timeperiod": 15})
    sma.set_input_arrays(k)
    k["SMA5"] = sma.outputs
    sma.set_parameters({"timeperiod": 20})
    k["SMA10"] = sma.outputs
    f = cross()
    print(f.run_score(k,dict(a="SMA5", b="SMA10")))